package com.google.code.mochaccino.framework.crypto.scheme.impl;

/*
 * Copyright 2012 Claude Houle claude.houle@gmail.com
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import com.google.code.mochaccino.framework.crypto.EncryptionException;
import com.google.code.mochaccino.framework.crypto.EncryptionProvider;
import com.google.code.mochaccino.framework.crypto.pgp.impl.PgpKeyProvider;
import com.google.code.mochaccino.framework.crypto.pgp.types.PgpCompression;
import com.google.code.mochaccino.framework.crypto.pgp.types.PgpLiteralDataMode;
import com.google.code.mochaccino.framework.crypto.pgp.types.PgpSymmetricKey;
import com.google.code.mochaccino.framework.crypto.scheme.EncryptionScheme;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.NoSuchProviderException;
import java.security.SecureRandom;
import java.util.Date;
import java.util.Iterator;
import org.apache.commons.io.IOUtils;
import org.bouncycastle.bcpg.ArmoredOutputStream;
import org.bouncycastle.openpgp.PGPCompressedData;
import org.bouncycastle.openpgp.PGPCompressedDataGenerator;
import org.bouncycastle.openpgp.PGPEncryptedDataGenerator;
import org.bouncycastle.openpgp.PGPEncryptedDataList;
import org.bouncycastle.openpgp.PGPException;
import org.bouncycastle.openpgp.PGPLiteralData;
import org.bouncycastle.openpgp.PGPLiteralDataGenerator;
import org.bouncycastle.openpgp.PGPObjectFactory;
import org.bouncycastle.openpgp.PGPOnePassSignatureList;
import org.bouncycastle.openpgp.PGPPublicKeyEncryptedData;
import org.bouncycastle.openpgp.PGPUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** PGP Implementation */
public class EncryptionSchemePGP implements EncryptionScheme {
	private Logger             logger                = LoggerFactory.getLogger( getClass() );
	private PgpCompression     compression           = null;
	private PgpKeyProvider     provider              = null;
	private PgpLiteralDataMode mode                  = PgpLiteralDataMode.BINARY;
	private PgpSymmetricKey    simmetricKeyAlgorithm = null;
	private SecureRandom       secureRandom          = new SecureRandom();

	public EncryptionSchemePGP( PgpKeyProvider keyProvider, PgpSymmetricKey algo, PgpCompression compression, PgpLiteralDataMode mode ) {
		this.simmetricKeyAlgorithm = algo;
		this.compression = compression;
		this.provider = keyProvider;
		this.mode = mode;
	}

	@Override
	public byte[] decrypt( byte[] payload ) throws EncryptionException {
		ByteArrayInputStream in = new ByteArrayInputStream(payload);
		ByteArrayOutputStream out = new ByteArrayOutputStream();
		decrypt( in, out );
		return out.toByteArray();
	}

	@Override
	public void decrypt( InputStream in, OutputStream out ) throws EncryptionException {
		InputStream decoderIn = null;
		try {
			decoderIn = PGPUtil.getDecoderStream( in );

			PGPObjectFactory decodedFactory = new PGPObjectFactory( decoderIn );
			Object decodedObject = null;
			while ( (decodedObject = decodedFactory.nextObject()) != null ) {
				if ( !(decodedObject instanceof PGPEncryptedDataList) ) {
					continue;
				}

				PGPEncryptedDataList dataList = (PGPEncryptedDataList) decodedObject;
				Iterator<?> dataIterator = dataList.getEncryptedDataObjects();
				while ( dataIterator.hasNext() ) {
					PGPPublicKeyEncryptedData data = (PGPPublicKeyEncryptedData) dataIterator.next();
					if ( data.isIntegrityProtected() && !data.verify() ) {
						throw new PGPException( "message failed integrity check" );
					}

					Object message = null;
					message = decrypt( data );
					message = decompress( message );

					if ( message instanceof PGPLiteralData ) {
						InputStream clearIn = null;
						try {
							PGPLiteralData clearData = (PGPLiteralData) message;
							clearIn = clearData.getInputStream();
							IOUtils.copy( clearIn, out );
						} finally {
							IOUtils.closeQuietly( clearIn );
						}
					} else if ( message instanceof PGPOnePassSignatureList ) {
						throw new PGPException( "encrypted message contains a signed message - not literal data." );
					} else {
						throw new PGPException( "message is not a simple encrypted file - type unknown." );
					}
				}
			}
		} catch ( PGPException e ) {
			throw new EncryptionException( e );
		} catch ( IOException e ) {
			throw new EncryptionException( e );
		} catch ( NoSuchProviderException e ) {
			throw new EncryptionException( e );
		} finally {
			IOUtils.closeQuietly( decoderIn );
			IOUtils.closeQuietly( in );
			IOUtils.closeQuietly( out );
		}
	}

	private Object decompress( Object message ) throws PGPException, IOException {
		if ( message instanceof PGPCompressedData ) {
			InputStream in = null;
			try {
				PGPCompressedData data = (PGPCompressedData) message;
				in = data.getDataStream();
				PGPObjectFactory factory = new PGPObjectFactory( in );
				message = factory.nextObject();
				return message;
			} finally {
				IOUtils.closeQuietly( in );
			}
		}
		return message;
	}

	private Object decrypt( PGPPublicKeyEncryptedData data ) throws PGPException, NoSuchProviderException, IOException {
		InputStream decryptedDataStream = data.getDataStream( provider.getPrivateKey(), EncryptionProvider.getBouncyCastleName() );
		PGPObjectFactory pgpDecryptedObjectFactory = new PGPObjectFactory( decryptedDataStream );
		return pgpDecryptedObjectFactory.nextObject();
	}

	@Override
	public byte[] encrypt( byte[] data ) throws EncryptionException {
		ByteArrayOutputStream out = new ByteArrayOutputStream();
		ByteArrayInputStream in = new ByteArrayInputStream( data );
		encrypt( in, out );
		return out.toByteArray();
	}

	@Override
	public void encrypt( InputStream in, OutputStream out ) throws EncryptionException {
		ArmoredOutputStream armoredOut = null;
		PGPEncryptedDataGenerator encryptor = null;
		OutputStream encryptedOutputStream = null;
		PGPCompressedDataGenerator compressor = null;
		OutputStream compressedOutputStream = null;
		PGPLiteralDataGenerator literalor = null;
		OutputStream literalOut = null;
		try {
			armoredOut = new ArmoredOutputStream( out );
			encryptor = new PGPEncryptedDataGenerator( simmetricKeyAlgorithm.getAlgo(), secureRandom, EncryptionProvider.getBouncyCastleName() );
			encryptor.addMethod( provider.getPublicKey() );
			encryptedOutputStream = encryptor.open( armoredOut, new byte[256] );
			compressor = new PGPCompressedDataGenerator( compression.getAlgo() );
			compressedOutputStream = compressor.open( encryptedOutputStream );
			literalor = new PGPLiteralDataGenerator();
			literalOut = literalor.open( compressedOutputStream, mode.getCode(), "internal", new Date(), new byte[256] );
			IOUtils.copy( in, literalOut );
		} catch ( IOException e ) {
			throw new EncryptionException( e );
		} catch ( PGPException e ) {
			throw new EncryptionException( e );
		} catch ( NoSuchProviderException e ) {
			throw new EncryptionException( e );
		} finally {
			IOUtils.closeQuietly( literalOut );
			closeQuietly( literalor );
			closeQuietly( compressor );
			IOUtils.closeQuietly( compressedOutputStream );
			closeQuietly( encryptor );
			IOUtils.closeQuietly( encryptedOutputStream );
			IOUtils.closeQuietly( armoredOut );
			IOUtils.closeQuietly( out );
			IOUtils.closeQuietly( in );
		}
	}

	private void closeQuietly( PGPEncryptedDataGenerator encryptor ) {
		try {
			if ( encryptor != null ) {
				encryptor.close();
			}
		} catch ( IOException e ) {
			logger.error( e.getMessage(), e );
		}
	}

	private void closeQuietly( PGPCompressedDataGenerator compressor ) {
		if ( compressor != null ) {
			try {
				compressor.close();
			} catch ( IOException e ) {
				logger.error( e.getMessage(), e );
			}
		}
	}

	private void closeQuietly( PGPLiteralDataGenerator literalor ) {
		try {
			if ( literalor != null ) {
				literalor.close();
			}
		} catch ( IOException e ) {
			logger.error( e.getMessage(), e );
		}
	}
}
